# -*- coding:utf-8 -*-

from ycdl import jijin as jj, line


# 业务处理
def do(ts_code, start_date, end_date, money, cycle, time, is_query):
    '''
        业务处理
    :param ts_code: 基金编号
    :param start_date: 开始时间
    :param end_date:   结束时间
    :param money:      投资金额
    :param cycle:      周期
    :param time:       时间
    :param is_query:   是否是查询收益，   False = 对比收益
    '''
    # 对日期作格式处理 去掉 ‘-’
    start_date = start_date.replace("-", '')
    end_date = end_date.replace('-', '')
    if is_query:
        # 查询收益
        obj = single(ts_code, start_date, end_date, money, cycle, time)
        return line.drawLine(obj, ts_code, start_date, end_date)
    else:
        # 对比
        obj, obj2 = compare(ts_code, start_date, end_date, money, cycle, time)
        return line.drawTwoLine(obj, obj2, ts_code, start_date, end_date)


# 查询收益
def single(ts_code, start_date, end_date, money, cycle, time):
    # 获取数据
    df = jj.get_data(ts_code, start_date, end_date)
    # 保存交易日期
    jj.trade_date = df.get('trade_date')
    # 保存收盘价
    jj.close = df.get('close')
    # 根据投资周期决定收益
    if cycle == 0:  # 按周
        # 每周定投
        obj = jj.weekModel(money, time)
    elif cycle == 1:
        # 每月定投
        obj = jj.monthModel(money, time)
    elif cycle == 2:
        # 每日定投
        obj = jj.dayModel(money)
    else:
        # 最优
        obj = jj.catchProfit(money, time)
    return obj


# 对比
def compare(ts_code, start_date, end_date, money, cycle, time):
    obj = single(ts_code, start_date, end_date, money, cycle, time)
    # 沪深 300
    ts_code2 = '510300.SH'
    obj2 = single(ts_code2, start_date, end_date, money, cycle, time)
    return obj, obj2
